//
// Created by LiuYou on 2021/10/22.
//

#include "MyTbf.hpp"
#include <cstdlib>
#include <unistd.h>
#include <algorithm>
#include <csignal>
#include <pthread.h>
#include <cstdio>
#include <cstring>


class MyTbf {
public:
    MyTbf() = default;

    ~MyTbf() = default;

public:
    int cps{};
    int burst{};
    int token{};
    int pos{};

public:
    /*
     * 加锁保护字段 token。
     */
    pthread_mutex_t mutex{};

};


/*
 * 这个策略机制数组对象, 是会造成线程竞争, 所以需要加锁保护。
 * 这个变量/对象是一个独占的资源使用。
 */
static MyTbf* tbfPolicyArray[MY_TBF_MAX];
static pthread_mutex_t MUTEX_TBF_POLICY_ARRAY = PTHREAD_MUTEX_INITIALIZER;

static pthread_t tidAlarm;

static pthread_once_t INIT_ONCE = PTHREAD_ONCE_INIT;


static int getFreePosUnlocked() {
    for (int i = 0; i < MY_TBF_MAX; ++i) {
        if (tbfPolicyArray[i] == nullptr) {
            return i;
        }
    }
    return -1;
}


/*
 * 把这个函数做成一个线程, 跑起来。
 * 这个线程, 每秒钟执行一次。
 * 可以用任何形式的 sleep();函数来实现。
 */
[[deprecated("在多线程重构令牌桶中已经弃用。")]]
static void alarmHandler(int sig) {
    for (int i = 0; i < MY_TBF_MAX; ++i) {
        if (tbfPolicyArray[i] != nullptr) {
            tbfPolicyArray[i]->token += tbfPolicyArray[i]->cps;
            if (tbfPolicyArray[i]->token > tbfPolicyArray[i]->burst) {
                tbfPolicyArray[i]->token = tbfPolicyArray[i]->burst;
            }
        }
    }
    alarm( 1 );
}


static void moduleUnload() {
    pthread_cancel( tidAlarm );

    for (int i = 0; i < MY_TBF_MAX; ++i) {
        if (tbfPolicyArray[i] != nullptr) {
            tbfDestroy( tbfPolicyArray[i] );
        }
    }


    pthread_mutex_destroy( &MUTEX_TBF_POLICY_ARRAY );

    pthread_join( tidAlarm, nullptr );
}


static void moduleLoad() {

    // 这里要创建一个线程, 这个线程干的事情就是 信号处理函数alarmHandler 要干的活。
    int errNum = pthread_create(
            &tidAlarm, nullptr,
            [](void* args) -> void* {

                while (true) {

                    pthread_testcancel();

                    pthread_mutex_lock( &MUTEX_TBF_POLICY_ARRAY );
                    for (int i = 0; i < MY_TBF_MAX; ++i) {
                        if (tbfPolicyArray[i] != nullptr) {

                            // 这两把锁互不影响。
                            pthread_mutex_lock( &tbfPolicyArray[i]->mutex );
                            tbfPolicyArray[i]->token += tbfPolicyArray[i]->cps;
                            if (tbfPolicyArray[i]->token > tbfPolicyArray[i]->burst) {
                                tbfPolicyArray[i]->token = tbfPolicyArray[i]->burst;
                            }
                            pthread_mutex_unlock( &tbfPolicyArray[i]->mutex );

                        }
                    }
                    pthread_mutex_unlock( &MUTEX_TBF_POLICY_ARRAY );

                    sleep( 1 );
                }

                pthread_exit( nullptr );
            },
            nullptr );
    if (errNum != 0) {
        fprintf( stderr, "pthread_create(): %s.\n", strerror( errNum ));
        // 这里就是要使用 exit();直接退出。
        // 因为如果每秒加令牌桶的线程都没有创建成功, 那还运行啥?
        exit( EXIT_FAILURE );
    }

    atexit( moduleUnload );
}


Tbf* tbfInit(int cps, int burst) {

    pthread_once( &INIT_ONCE, moduleLoad );

    //region 从临界区代码中 提取的与 独占资源无关的代码, 放在临界区外面。
    MyTbf* anInstance = static_cast<MyTbf*>(malloc( sizeof( MyTbf )));
    if (anInstance == nullptr) {
        return nullptr;
    }

    anInstance->token = 0;
    anInstance->cps = cps;
    anInstance->burst = burst;

    // 在 malloc();之后对锁进行初始化。如果是C++就可以使用 RAII机制。
    pthread_mutex_init( &anInstance->mutex, nullptr );
    //endregion


    pthread_mutex_lock( &MUTEX_TBF_POLICY_ARRAY );
    int pos = getFreePosUnlocked();
    if (pos < 0) {
        pthread_mutex_unlock( &MUTEX_TBF_POLICY_ARRAY );

        free( anInstance );
        return nullptr;
    }


    anInstance->pos = pos;

    tbfPolicyArray[pos] = anInstance;
    pthread_mutex_unlock( &MUTEX_TBF_POLICY_ARRAY );


    return anInstance;
}

int tbfDestroy(Tbf* tbf) {
    MyTbf* anInstance = static_cast<MyTbf*>(tbf);

    pthread_mutex_lock( &MUTEX_TBF_POLICY_ARRAY );
    tbfPolicyArray[anInstance->pos] = nullptr;
    pthread_mutex_unlock( &MUTEX_TBF_POLICY_ARRAY );

    // 如果是C就要这样写, C++可以使用 RAII机制。
    pthread_mutex_destroy( &anInstance->mutex );
    free( tbf );
    return 0;
}

int fetchToken(Tbf* tbf, int tokens) {
    if (tokens <= 0) {
        return -1;
    }

    MyTbf* anInstance = static_cast<MyTbf*>(tbf);

    pthread_mutex_lock( &anInstance->mutex );
    while (anInstance->token <= 0) {
        pthread_mutex_unlock( &anInstance->mutex );
        sched_yield();
        pthread_mutex_lock( &anInstance->mutex );
    }

    int min = std::min( anInstance->token, tokens );
    anInstance->token -= min;
    pthread_mutex_unlock( &anInstance->mutex );

    return min;
}

int returnToken(Tbf* tbf, int tokens) {
    if (tokens <= 0) {
        return -1;
    }

    MyTbf* me = static_cast<MyTbf*>(tbf);

    pthread_mutex_lock( &me->mutex );
    me->token += tokens;
    if (me->token > me->burst) {
        me->token = me->burst;
    }
    pthread_mutex_unlock( &me->mutex );


    return tokens;
}
